use chrono::Local;
use walkdir::WalkDir;
use yara::Compiler;

// Note: this requires the `derive` feature

use std::path::PathBuf;
use std::ffi::OsString;
use std::os::unix::net::UnixDatagram;

use clap::{Args, Parser, Subcommand};

pub mod model;

/// Tailmon-EDR scan CLI
#[derive(Debug, Parser)]
#[clap(name = "tailmonscan")]
#[clap(about = "Tailmon-EDR scan CLI", long_about = None)]
struct Cli {
    #[clap(subcommand)]
    command: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    /// Clones repos
    #[clap(arg_required_else_help = true)]
    Clone {
        /// The remote to clone
        remote: String,
    },
    /// pushes things
    #[clap(arg_required_else_help = true)]
    Push {
        /// The remote to target
        remote: String,
    },
    /// adds things
    #[clap(arg_required_else_help = true)]
    Add {
        /// Stuff to add
        #[clap(required = true, parse(from_os_str))]
        path: Vec<PathBuf>,
    },
    /// Virus scan
    #[clap(arg_required_else_help = true)]
    Virus {
        /// Rule path
        #[clap(short, long, parse(from_os_str))]
        rule: Option<PathBuf>,
        /// Threads number
        #[clap(short, long)]
        threads: Option<i32>,
        /// Enable process scan
        #[clap(short, long, takes_value = true, default_missing_value = "true")]
        process: Option<bool>,
        /// Scan path
        #[clap(required = true, parse(from_os_str))]
        path: Vec<PathBuf>,
    },
    Stash(Stash),
    #[clap(external_subcommand)]
    External(Vec<OsString>),
}

#[derive(Debug, Args)]
#[clap(args_conflicts_with_subcommands = true)]
struct Stash {
    #[clap(subcommand)]
    command: Option<StashCommands>,

    #[clap(flatten)]
    push: StashPush,
}

#[derive(Debug, Subcommand)]
enum StashCommands {
    Push(StashPush),
    Pop { stash: Option<String> },
    Apply { stash: Option<String> },
}

#[derive(Debug, Args)]
struct StashPush {
    #[clap(short, long)]
    message: Option<String>,
}

fn main() {
    let args = Cli::parse();

    match args.command {
        Commands::Clone { remote } => {
            println!("Cloning {}", remote);
        }
        Commands::Push { remote } => {
            println!("Pushing to {}", remote);
        }
        Commands::Add { path } => {
            println!("Adding {:?}", path);
        }
        Commands::Virus {
            rule,
            threads,
            process,
            path,
        } => {
            println!("Virus Scan Starting...");
            let result = virus_scan(rule, threads, path, process);
            if result.is_ok() {
                println!("Virus Scan Finished!");
            } else {
                println!("{:?}", result);
            }
            
        }
        Commands::Stash(stash) => {
            let stash_cmd = stash.command.unwrap_or(StashCommands::Push(stash.push));
            match stash_cmd {
                StashCommands::Push(push) => {
                    println!("Pushing {:?}", push);
                }
                StashCommands::Pop { stash } => {
                    println!("Popping {:?}", stash);
                }
                StashCommands::Apply { stash } => {
                    println!("Applying {:?}", stash);
                }
            }
        }
        Commands::External(args) => {
            println!("Calling out to {:?} with {:?}", &args[0], &args[1..]);
        }
    }

    // Continued program logic goes here...
}

fn virus_scan(rule: Option<PathBuf>, threads: Option<i32>, path: Vec<PathBuf>, process_scan: Option<bool>) -> std::io::Result<()> {
    let mut task_total: i64 = 0;
	let mut task_finish: i64 = 0;
	let mut task_match: i64 = 0;

    // 统计总文件数
    for scan_path in path.clone() {
        for file in WalkDir::new(scan_path)
            .into_iter()
            .filter_map(|file| file.ok()) {
            if file.metadata().unwrap().is_file() {
                task_total += 1;
            }
        }
    }
    // 统计总进程数
    if process_scan.unwrap_or(false) {
        let processes = psutil::process::processes().unwrap();
        for p in processes {
            if p.unwrap().pid() > 10000 {
                task_total += 1;
            }
            
        }
    }

    // init UnixDatagram
    let socket = UnixDatagram::unbound()?;

    // init Compiler
    let rule_path = rule.unwrap_or(PathBuf::from(
        r"/root/projects/tailmon-edr/tailmond/data/rules/current.yara",
    ));
    let _threads_number = threads.unwrap_or(1);

    let compiler = Compiler::new().unwrap();
    let compiler = compiler
        .add_rules_file(rule_path)
        .expect("Should have parsed rule");

    let rules = compiler
        .compile_rules()
        .expect("Should have compiled rules");

    println!("File scanning...");
    for scan_path in path {
        for file in WalkDir::new(scan_path)
            .into_iter()
            .filter_map(|file| file.ok())
        {
            if file.metadata().unwrap().is_file() {
                
                let result = rules.scan_file(file.path(), 30);
                task_finish += 1;
                if result.is_ok() {
                    
                    let timestamp = Local::now().timestamp_millis();
                    let rule_id: Vec<String> = result.unwrap().iter().map(|r| r.identifier.to_string()).collect();
                    let mut matches = false;
                    if !rule_id.is_empty() {
                        task_match += 1;
                        matches = true;
                        println!("{} matches:{:?}", file.path().display(), rule_id);
                    } else {
                        println!("{} no matches", file.path().display());
                    }
                    let yara_result = model::YaraResult {
                        matches,
                        timestamp,
                        mal_obj_type: "File".to_string(),
                        mal_obj_key: file.path().display().to_string(),
                        rule_id,
                        task_total,
                        task_finish,
                        task_match,
                    };
                    let yara_result_json = serde_json::to_string(&yara_result).unwrap() + "\n";

                    socket.send_to(yara_result_json.as_bytes(), "/var/run/yara.sock")?;
                } else {
                    println!("{:?}", result);
                }

            }
        }
    }

    println!("{:?}", process_scan);
    if process_scan.unwrap_or(false) {
        println!("Process scanning...");
        // 进程扫描
        let processes = psutil::process::processes().unwrap();
        for p in processes {
            let p = p.unwrap();
            let pid = p.pid();
            if pid <= 10000 {
                continue;
            }

            let result = rules.scan_process(pid, 600);
                task_finish += 1;
                if result.is_ok() {
                    
                    let timestamp = Local::now().timestamp_millis();
                    let rule_id: Vec<String> = result.unwrap().iter().map(|r| r.identifier.to_string()).collect();
                    let mut matches = false;
                    if !rule_id.is_empty() {
                        task_match += 1;
                        matches = true;
                        println!("PID {} matches:{:?}", pid, rule_id);
                    } else {
                        println!("PID {} no matches", pid);
                    }
                    let yara_result = model::YaraResult {
                        matches,
                        timestamp,
                        mal_obj_type: "Process".to_string(),
                        mal_obj_key: pid.to_string(),
                        rule_id,
                        task_total,
                        task_finish,
                        task_match,
                    };
                    let yara_result_json = serde_json::to_string(&yara_result).unwrap() + "\n";

                    socket.send_to(yara_result_json.as_bytes(), "/var/run/yara.sock")?;
                } else {
                    println!("{:?}", result);
                }
        }
    }

    
    Ok(())
}
